const Yolov5Detection = require("./yolov5Detection.js");
const Yolov8Detection = require("./yolov8Detection.js");
const fs = require("fs");
const path = require("path");
const ort = require("onnxruntime-node");

async function yolov5(imagePath, confidenceThreshold) {
  const yolov5Detection = new Yolov5Detection();
  await yolov5Detection.createModel("./onnx/yolov5n.onnx");
  const detections = await yolov5Detection.detect(
    imagePath,
    640,
    confidenceThreshold,
    0.5
  );
  console.log(detections);

  const outputDir = "./output/yolov5";
  if (!fs.existsSync(outputDir)) {
    fs.mkdirSync(outputDir, { recursive: true });
  }
  const outputFileName = `detected_${path.basename(imagePath)}`;
  const outputPath = path.join(outputDir, outputFileName);

  yolov5Detection.drawBoundingBoxes(
    imagePath,
    detections,
    outputPath,
    640,
    "red",
    2,
    "green",
    24
  );
}

async function yolov8(imagePath, confidenceThreshold) {
  const yolov8Detection = new Yolov8Detection();
  await yolov8Detection.createModel("./onnx/yolov8n_nms_simplify.onnx");
  const detections = await yolov8Detection.detect(
    imagePath,
    640,
    confidenceThreshold,
    0.5
  );
  console.log(detections);

  const outputDir = "./output/yolov8";
  if (!fs.existsSync(outputDir)) {
    fs.mkdirSync(outputDir, { recursive: true });
  }
  const outputFileName = `detected_${path.basename(imagePath)}`;
  const outputPath = path.join(outputDir, outputFileName);

  yolov8Detection.drawBoundingBoxes(
    imagePath,
    detections,
    outputPath,
    640,
    "red",
    2,
    "green",
    24
  );
}
const imagePath = "./images/scratches_64.jpg";
const confidenceThreshold = 0.5;
yolov5(imagePath, confidenceThreshold);
yolov8(imagePath, confidenceThreshold);
